/*
 * Copyright 2010-2025 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlin.cli.pipeline

import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.backend.utils.createSyntheticFirFileForFir2Ir
import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
import org.jetbrains.kotlin.fir.declarations.FirFile
import org.jetbrains.kotlin.fir.declarations.utils.fileNameForPluginGeneratedCallable
import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
import org.jetbrains.kotlin.fir.extensions.declarationGenerators
import org.jetbrains.kotlin.fir.extensions.extensionService
import org.jetbrains.kotlin.fir.extensions.generatedDeclarationsSymbolProvider
import org.jetbrains.kotlin.fir.moduleData
import org.jetbrains.kotlin.fir.pipeline.AllModulesFrontendOutput
import org.jetbrains.kotlin.fir.symbols.SymbolInternals
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.packageName

/**
 * This phase generates synthetic files for top-level declarations, generated by compiler plugins.
 * Having this step as a separate phase allows consistently using this logic in all CLI pipelines
 * including the Metadata compiler.
 */
class FrontendFilesForPluginsGenerationPipelinePhase<A : FrontendPipelineArtifact> : PipelinePhase<A, A>(
    name = "FrontendFilesForPluginsGenerationPipelinePhase"
) {
    override fun executePhase(input: A): A {
        return input.withNewFrontendOutput(
            createFilesWithGeneratedDeclarations(input.frontendOutput)
        )
    }

    companion object {
        fun createFilesWithGeneratedDeclarations(allModulesOutput: AllModulesFrontendOutput): AllModulesFrontendOutput {
            val outputs = allModulesOutput.outputs.map {
                val newFiles = createFilesWithGeneratedDeclarations(it.session)
                if (newFiles.isEmpty()) return@map it
                it.copy(fir = it.fir + newFiles)
            }
            return AllModulesFrontendOutput(outputs)
        }

        @OptIn(FirExtensionApiInternals::class, SymbolInternals::class)
        fun createFilesWithGeneratedDeclarations(session: FirSession): List<FirFile> {
            val symbolProvider = session.generatedDeclarationsSymbolProvider ?: return emptyList()
            val declarationGenerators = session.extensionService.declarationGenerators

            val fileModuleData = session.moduleData

            return buildList {
                val generatedClasses = declarationGenerators.flatMap { it.topLevelClassIdsCache.getValue() }
                    .mapNotNull { symbolProvider.getClassLikeSymbolByClassId(it)?.fir }

                // classes go to a specific file each
                for (firClass in generatedClasses) {
                    val classId = firClass.symbol.classId
                    this += createSyntheticFirFileForFir2Ir(
                        fileName = "${classId.packageFqName.toPath()}/${classId.relativeClassName.asString()}.kt",
                        packageFqName = classId.packageFqName,
                        fileModuleData,
                        FirDeclarationOrigin.Synthetic.PluginFile,
                    ) {
                        declarations += firClass
                    }
                }

                // callables are grouped per package
                val generatedCallables = declarationGenerators.flatMap { it.topLevelCallableIdsCache.getValue() }
                    .flatMap { symbolProvider.getTopLevelCallableSymbols(it.packageName, it.callableName) }

                val generatedCallablesPerPackage = generatedCallables.groupBy { it.callableId.packageName }
                for ((packageName, packageGeneratedCallables) in generatedCallablesPerPackage) {
                    val callablesPerFileName = packageGeneratedCallables.groupBy {
                        val name = it.fir.fileNameForPluginGeneratedCallable ?: "__GENERATED__CALLABLES__.kt"
                        if (name.endsWith(".kt")) name else "$name.kt"
                    }
                    for ((fileName, callables) in callablesPerFileName) {
                        this += createSyntheticFirFileForFir2Ir(
                            fileName = "${packageName.toPath()}/$fileName",
                            packageFqName = packageName,
                            fileModuleData,
                            FirDeclarationOrigin.Synthetic.PluginFile,
                        ) {
                            declarations += callables.map { it.fir }
                        }
                    }
                }
            }
        }

        private fun FqName.toPath(): String = this.asString().replace('.', '/')
    }
}
